package github.sf.fw.tools.db.sql;

import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.expr.SQLIdentifierExpr;
import com.alibaba.druid.sql.ast.expr.SQLPropertyExpr;
import com.alibaba.druid.sql.ast.statement.*;
import com.alibaba.druid.sql.visitor.SQLASTVisitorAdapter;

import java.util.List;

public class ModifyTableSourceVisitor extends SQLASTVisitorAdapter {

    private String database = "JwCloud";
    private int modifyCount = 0;

    public void setDatabase(String database) {
        this.database = database;
    }

    public boolean visit(SQLLateralViewTableSource x) {
        modifyCount++;
        return true;
    }

    public boolean visit(SQLValuesTableSource x) {
        modifyCount++;
        return true;
    }

    @Override
    public void endVisit(SQLValuesTableSource x) {

    }

    public boolean visit(SQLUnionQueryTableSource x) {
        modifyCount++;
        return true;
    }

    @Override
    public void endVisit(SQLUnionQueryTableSource x) {

    }

    public boolean visit(SQLSubqueryTableSource x) {
        modifyCount++;
        return true;
    }

    @Override
    public void endVisit(SQLSubqueryTableSource x) {

    }

    public boolean visit(SQLJoinTableSource x) {
        modifyCount++;
        return true;
    }

    @Override
    public void endVisit(SQLJoinTableSource x) {

    }

    public boolean visit(SQLWithSubqueryClause.Entry x) {
        modifyCount++;
        return true;
    }

    @Override
    public void endVisit(SQLWithSubqueryClause x) {
        List<SQLWithSubqueryClause.Entry> entries = x.getEntries();
        for (SQLWithSubqueryClause.Entry entry : entries) {
            System.out.println(entry);
        }
    }

    public boolean visit(SQLExprTableSource x) {
        modifyCount++;
        return true;
    }

    @Override
    public void endVisit(SQLExprTableSource x) {
        modifyTable(x);
    }

    protected void modifyTable(SQLExprTableSource x) {
        SQLExpr expr = x.getExpr();
        if (expr instanceof SQLIdentifierExpr) {
            SQLPropertyExpr sqlPropertyExpr = new SQLPropertyExpr();
            sqlPropertyExpr.setName(((SQLIdentifierExpr) expr).getName());
            SQLIdentifierExpr sqlIdentifierExpr = new SQLIdentifierExpr();
            sqlIdentifierExpr.setName(database);
            sqlPropertyExpr.setOwner(sqlIdentifierExpr);
            x.setExpr(sqlPropertyExpr);
        } else if (expr instanceof SQLPropertyExpr) {
            SQLExpr owner = ((SQLPropertyExpr) expr).getOwner();
            if (owner instanceof SQLIdentifierExpr) {
                ((SQLIdentifierExpr) owner).setName(database);
            }
        }
    }

    public int getModifyCount() {
        return modifyCount;
    }
}
